import torch
import torch.nn as nn
from matplotlib import pyplot as plt
import numpy as np

# 设置字体为支持中文的字体，例如 SimHei（黑体）
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置字体为 SimHei
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载weight
weight_focus = torch.load('weight_focus_tensor.pt', map_location=device)
print(f"weight focus={weight_focus}")

# # # 加载原始高分辨率复数核
kernel_focus = torch.load('kernel_focus_tensor.pt', map_location=device)
#
# def fft2(data):
#     data = torch.fft.ifftshift(data, dim=(-2, -1))
#     data = torch.fft.fftn(data, dim=[-1, -2])
#     data = torch.fft.fftshift(data, dim=(-2, -1))
#     return data
#
#
# def ifft2(data):
#     data = torch.fft.ifftshift(data, dim=(-2, -1))
#     data = torch.fft.ifftn(data, dim=[-1, -2])
#     data = torch.fft.fftshift(data, dim=(-2, -1))
#     return data

# 挨个取出核的abs值
# for i in range(24):
#     kernel_focus_abs = torch.abs(kernel_focus[i])
#     plt.imshow(kernel_focus_abs.cpu().numpy())
#     plt.title(f'高分辨率第{i+1}个核绝对值')
#     plt.colorbar()
#     plt.show()

# 挨个绘制核的频谱
# for i in range(24):
#     kernel_focus_fft = fft2(kernel_focus[i])
#     kernel_focus_abs = torch.abs(kernel_focus_fft)
#     plt.imshow(kernel_focus_abs.cpu().numpy(), origin='lower', cmap='gray')
#     plt.title(f'高分辨率第{i+1}个核频谱绝对值')
#     plt.colorbar()
#     plt.show()

# 挨个绘制核的实部虚部
# for i in range(24):
#     kernel_focus_i = kernel_focus[i]
#     kernel_focus_real = kernel_focus_i.real.cpu().numpy()
#     kernel_focus_imag = kernel_focus_i.imag.cpu().numpy()
#     plt.imshow(kernel_focus_real)
#     plt.title(f'高分辨率第{i+1}个核实部')
#     plt.colorbar()
#     plt.show()
#     plt.imshow(kernel_focus_imag)
#     plt.title(f'高分辨率第{i + 1}个核虚部')
#     plt.colorbar()
#     plt.show()

# for i in range(24):
#     kernel_focus_i = kernel_focus[i]
#     kernel_focus_i_ifft = torch.fft.ifft2(torch.fft.ifftshift(kernel_focus_i))
#     kernel_focus_real = kernel_focus_i_ifft.real.cpu().numpy()
#     kernel_focus_imag = kernel_focus_i_ifft.imag.cpu().numpy()
#     plt.imshow(kernel_focus_real)
#     plt.title(f'高分辨率第{i+1}个核实部')
#     plt.colorbar()
#     plt.show()
#     plt.imshow(kernel_focus_imag)
#     plt.title(f'高分辨率第{i + 1}个核虚部')
#     plt.colorbar()
#     plt.show()


# for i in range(24):
#     N = 35  # 数组大小
#     # 计算二维傅里叶变换
#     fft_signal = torch.fft.fft2(kernel_focus[i])
#     # 计算频谱的幅度
#     magnitude = torch.abs(torch.fft.ifftshift(fft_signal))
#     # 获取频率轴
#     frequencies_x = torch.fft.fftfreq(N, 1 / N)  # X轴的频率
#     frequencies_y = torch.fft.fftfreq(N, 1 / N)  # Y轴的频率
#     plt.imshow(magnitude.cpu().numpy(), origin='lower',cmp='jet')
#     plt.title(f'高分辨率第{i+1}个核频谱')
#     plt.colorbar()
#     plt.show()


# for i in range(24):
#     kernel_focus = kernel_focus[i]
#     plt.imshow(.cpu().numpy())
#     plt.title(f'高分辨率第{i+1}个核绝对值')
#     plt.colorbar()
#     plt.show()

# # 加载原始高分辨率共轭复数核,共轭核abs与原始相同 a+bi, a-bi
# kernel_focus_ct = torch.load('kernel_ct_focus_tensor.pt', map_location=device)
#
# for i in range(1):
#     kernel_focus_ct_real = kernel_focus_ct[i].real
#     kernel_focus_ct_real = kernel_focus_ct[i].imag
#     plt.imshow(kernel_focus_ct_real.cpu().numpy())
#     plt.title(f'高分辨率第{i+1}个共轭核绝对值')
#     plt.colorbar()
#     plt.show()
#

# # 发现降采样kernel做法是正确的，但计算不对
# kernel_focus_lowres4 = torch.load('kernel_focus_tensor_lowres4.pt', map_location=device)
# for i in range(3):
#     kernel_focus_abs = torch.abs(kernel_focus_lowres4[i])
#     plt.imshow(kernel_focus_abs.cpu().numpy())
#     plt.title(f'降低分辨率4倍第{i+1}个核绝对值')
#     plt.colorbar()
#     plt.show()
#
# kernel_focus_lowres2 = torch.load('kernel_focus_tensor_lowres2.pt', map_location=device)
# for i in range(3):
#     kernel_focus_abs = torch.abs(kernel_focus_lowres2[i])
#     plt.imshow(kernel_focus_abs.cpu().numpy())
#     plt.title(f'降低分辨率2倍第{i+1}个核绝对值')
#     plt.colorbar()
#     plt.show()

#
# plt.imshow(kernel_focus_lowres2_real_mean.cpu().numpy())
# plt.title('2倍降采样核实数部分均值')
# plt.colorbar()
# plt.show()
# plt.imshow(kernel_focus_lowres2_imag_mean.cpu().numpy())
# plt.title('2倍降采样核虚数部分均值')
# plt.colorbar()
# plt.show()


# # 定义降采样层
# avg_layer = nn.AvgPool2d(kernel_size=2, stride=2)
#
# # 检查是否为复数张量
# if torch.is_complex(kernel_focus):
#     # 分离实部和虚部
#     kernel_real = kernel_focus.real
#     kernel_imag = kernel_focus.imag
#
#     # 对实部和虚部分别降采样
#     kernel_real_lowres = avg_layer(kernel_real.unsqueeze(1)).squeeze(1)  # [24, 512, 512]
#     kernel_imag_lowres = avg_layer(kernel_imag.unsqueeze(1)).squeeze(1)  # [24, 512, 512]
#
#     # 合并为复数核
#     kernel_focus_lowres = torch.complex(kernel_real_lowres, kernel_imag_lowres)
# else:
#     # 如果是实数核，直接降采样
#     kernel_focus_lowres = avg_layer(kernel_focus.unsqueeze(1)).squeeze(1)
#
# # 保存低分辨率核
# torch.save(kernel_focus_lowres, 'kernel_ct_defocus_tensor_lowres2.pt')